package com.zsc.oj.manager;

import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;

import com.zsc.oj.common.ErrorCode;
import com.zsc.oj.exception.BusinessException;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Component
public class AIManager {

    @Resource
    private ClientV4 clientV4;

    public static final Float STABLE_TEMPERATURE = 0.05f;

    public static final Float UNSTABLE_TEMPERATURE = 0.99F;

    /**
     * 回答不稳定的请求
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String doUnstableRequest(String systemMessage,String userMessage){
        return doRequest(systemMessage,userMessage,Boolean.FALSE,UNSTABLE_TEMPERATURE);
    }

    /**
     * 回答稳定的请求
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String doStableRequest(String systemMessage,String userMessage){
        return doRequest(systemMessage,userMessage,Boolean.FALSE,STABLE_TEMPERATURE);
    }

    /**
     * 同步请求
     * @param messages
     * @param stream
     * @param temperature
     * @return
     */
    public String doSyncRequest(List<ChatMessage> messages,Boolean stream,Float temperature){
        return doRequest(messages, Boolean.FALSE, temperature);
    }

    /**
     * 提问信息响应
     * @param systemMessage：系统prompt
     * @param stream：用户prompt
     * @param temperature
     * @return
     */
    public String doRequest(String systemMessage,String userMessage,Boolean stream,Float temperature){
        //构建请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(),systemMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(),userMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doRequest(messages,stream,temperature);
    }


    /**
     * 全局通用响应
     * @param messages
     * @param stream
     * @param temperature
     * @return
     */
    public String doRequest(List<ChatMessage> messages,Boolean stream,Float temperature){
        //构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                //模型选择
                .model(Constants.ModelChatGLM4)
                .stream(stream)
                .invokeMethod(Constants.invokeMethod)
                .temperature(temperature)
                .messages(messages)
                .build();
        try {
            ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
            return invokeModelApiResp.getData().getChoices().get(0).toString();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, e.getMessage());
        }
    }

    /**
     * 流式回答不稳定的请求
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public Flowable<ModelData> doStreamUnstableRequest(String systemMessage,String userMessage){
        return doStreamRequest(systemMessage,userMessage,UNSTABLE_TEMPERATURE);
    }

    /**
     * 流式回答稳定的请求
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public Flowable<ModelData> doStreamStableRequest(String systemMessage,String userMessage){
        return doStreamRequest(systemMessage,userMessage,STABLE_TEMPERATURE);
    }

    /**
     * 提问信息流式响应
     * @param systemMessage：系统prompt
     * @param userMessage：用户prompt
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doStreamRequest(String systemMessage,String userMessage,Float temperature){
        //构建请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(),systemMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(),userMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doRequest(messages,temperature);
    }


    /**
     * 流式通用响应
     * @param messages
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doRequest(List<ChatMessage> messages, Float temperature){
        //构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                //模型选择
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .invokeMethod(Constants.invokeMethod)
                .temperature(temperature)
                .messages(messages)
                .build();
        try {
            ModelApiResponse invokeModelApiResp = clientV4.invokeModelApi(chatCompletionRequest);
            return invokeModelApiResp.getFlowable();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, e.getMessage());
        }
    }

}
